# SPDX-License-Identifier: MIT
# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.

import functools
import inspect
import logging


def execute_operation(operation_str, param_dict):
    if operation_str:
        try:
            exec(operation_str, param_dict)
        except Exception as e:
            logging.error(f"Error executing exit code: {e}")


def marker_prof_wrapper(original_method, params):
    function_name = params.get("function_name", None)
    if function_name is None:
        return original_method

    entry_operation = params.get("entry_operation", None)
    exit_operation = params.get("exit_operation", None)
    scope_name = params.get("scope_name", None)
    scope_class_name = params.get("scope_class_name", None)
    scope_function = params.get("scope_function", None)

    if inspect.iscoroutinefunction(original_method):
        logging.info(f"<<<INFO: {original_method.__qualname__} is async function, use async wrapper")

        @functools.wraps(original_method)
        async def async_wrapper(self, *args, **kwargs):
            entry_args = args
            entry_kwargs = kwargs
            param_dict = {"self": self, "args": entry_args, "kwargs": entry_kwargs}
            execute_operation(None, {})

            caller_module = ""
            caller_class = ""
            caller_function = ""

            # 如果未指定范围，总是执行操作
            if not all([scope_name, scope_class_name, scope_function]):
                should_wrap = True
            else:
                # 只检查直接调用者
                should_wrap = False
                try:
                    # 获取当前调用栈
                    stack = inspect.stack()
                    if len(stack) > 1:  # 确保有调用者
                        # 直接调用者是栈的第二帧（索引1）
                        caller_frame = stack[1].frame

                        # 获取调用者的模块和类信息
                        caller_module = caller_frame.f_globals.get('__name__', '')

                        # 获取调用者的类名（如果有self参数）
                        caller_self = caller_frame.f_locals.get('self', None)
                        caller_class = caller_self.__class__.__name__ if caller_self else ''

                        # 获取调用函数名
                        caller_function = stack[1].function

                        # 检查是否匹配
                        should_wrap = (caller_module == scope_name and
                                       caller_class == scope_class_name and
                                       caller_function == scope_function)
                except Exception as e:
                    logging.error(f"Error checking direct caller: {e}")
                    should_wrap = False

            # 如果不需要包装，直接返回原始方法
            if not should_wrap:
                logging.info(
                    f"<<<INFO: Skipping wrapper for {original_method.__qualname__} due to scope {caller_module}.{caller_class}.{caller_function} mismatch")
                return original_method(self, *args, **kwargs)

            execute_operation(entry_operation, param_dict)
            result = await original_method(self, *args, **kwargs)
            param_dict["result"] = result
            execute_operation(exit_operation, param_dict)
            return result

        return async_wrapper
    else:
        logging.info(f"<<<INFO: {original_method.__qualname__} is sync function, use sync wrapper")

        @functools.wraps(original_method)
        def wrapper(self, *args, **kwargs):

            entry_args = args
            entry_kwargs = kwargs
            param_dict = {"self": self, "args": entry_args, "kwargs": entry_kwargs}
            execute_operation(None, {})

            caller_module = ""
            caller_class = ""
            caller_function = ""

            # 如果未指定范围，总是执行操作
            if not all([scope_name, scope_class_name, scope_function]):
                should_wrap = True
            else:
                # 只检查直接调用者
                should_wrap = False
                try:
                    # 获取当前调用栈
                    stack = inspect.stack()
                    if len(stack) > 1:  # 确保有调用者
                        # 直接调用者是栈的第二帧（索引1）
                        caller_frame = stack[1].frame

                        # 获取调用者的模块和类信息
                        caller_module = caller_frame.f_globals.get('__name__', '')

                        # 获取调用者的类名（如果有self参数）
                        caller_self = caller_frame.f_locals.get('self', None)
                        caller_class = caller_self.__class__.__name__ if caller_self else ''

                        # 获取调用函数名
                        caller_function = stack[1].function

                        # 检查是否匹配
                        should_wrap = (caller_module == scope_name and
                                       caller_class == scope_class_name and
                                       caller_function == scope_function)
                except Exception as e:
                    logging.error(f"Error checking direct caller: {e}")
                    should_wrap = False

            # 如果不需要包装，直接返回原始方法
            if not should_wrap:
                logging.info(
                    f"<<<INFO: Skipping wrapper for {original_method.__qualname__} due to scope {caller_module}.{caller_class}.{caller_function} mismatch")
                return original_method(self, *args, **kwargs)

            execute_operation(entry_operation, param_dict)
            result = original_method(self, *args, **kwargs)
            param_dict["result"] = result
            execute_operation(exit_operation, param_dict)

            return result

        return wrapper
